"""
Articulation subset management for robotics simulation.
"""

import numpy as np
from isaacsim.core.utils.types import ArticulationAction


class ArticulationSubset:
    """简化的关节子集实现，替代internutopia.core.robot.articulation_subset"""

    def __init__(self, articulation, joint_names):
        self.articulation = articulation
        self.joint_names = joint_names
        self._joint_indices = None
        self._initialize_joint_indices()

    def _initialize_joint_indices(self):
        """初始化关节索引映射"""
        if not hasattr(self.articulation, 'handles_initialized') or not self.articulation.handles_initialized:
            print("Warning: Articulation handles not initialized")
            self._joint_indices = []
            return

        self._joint_indices = []
        for joint_name in self.joint_names:
            try:
                index = self.articulation.get_dof_index(joint_name)
                self._joint_indices.append(index)
            except Exception as e:
                print(f"Warning: Could not find joint '{joint_name}': {e}")
                self._joint_indices.append(-1)

        print(f"Initialized {len(self.joint_names)} joint indices")

    def get_joint_positions(self):
        """获取子集关节位置"""
        positions = self.articulation.get_joint_positions()
        if len(positions.shape) == 1:
            return positions[self._joint_indices]
        else:
            return positions[:, self._joint_indices]

    def get_joint_velocities(self):
        """获取子集关节速度"""
        velocities = self.articulation.get_joint_velocities()
        if len(velocities.shape) == 1:
            return velocities[self._joint_indices]
        else:
            return velocities[:, self._joint_indices]

    def set_joint_positions(self, positions):
        """设置子集关节位置"""
        self.articulation.set_joint_positions(positions, joint_indices=self._joint_indices)

    def set_joint_velocities(self, velocities):
        """设置子集关节速度"""
        self.articulation.set_joint_velocities(velocities, joint_indices=self._joint_indices)

    def make_articulation_action(self, joint_positions=None, joint_velocities=None):
        """创建仅针对子集关节的ArticulationAction"""
        return ArticulationAction(
            joint_positions=joint_positions,
            joint_velocities=joint_velocities,
            joint_indices=self._joint_indices
        )

    def apply_action(self, joint_positions=None, joint_velocities=None):
        """应用动作到子集关节"""
        action = self.make_articulation_action(joint_positions, joint_velocities)
        self.articulation.apply_action(action)

    def get_joint_subset_indices(self):
        """获取子集关节在完整articulation中的索引"""
        return self._joint_indices